import asyncio
import json
import logging
import websockets
from typing import Callable, Optional, Dict, Any
from datetime import datetime
import threading
import time


class WebSocketClient:
    """WebSocket客户端，连接到C++后端服务器"""
    
    def __init__(self, uri: str = "ws://localhost:8080"):
        self.uri = uri
        self.websocket: Optional[websockets.WebSocketServerProtocol] = None
        self.connected = False
        self.running = False
        self.message_handlers: Dict[str, Callable] = {}
        self.connection_handlers: list[Callable] = []
        self.disconnection_handlers: list[Callable] = []
        self.logger = logging.getLogger(__name__)
        self.reconnect_interval = 5  # 重连间隔（秒）
        self.max_reconnect_attempts = 10
        self.reconnect_attempts = 0
        
    async def connect(self) -> bool:
        """连接到WebSocket服务器"""
        try:
            self.websocket = await websockets.connect(self.uri)
            self.connected = True
            self.reconnect_attempts = 0
            self.logger.info(f"已连接到WebSocket服务器: {self.uri}")
            
            # 通知连接处理器
            for handler in self.connection_handlers:
                try:
                    handler()
                except Exception as e:
                    self.logger.error(f"连接处理器异常: {e}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"连接WebSocket服务器失败: {e}")
            self.connected = False
            return False
    
    async def disconnect(self):
        """断开WebSocket连接"""
        self.running = False
        self.connected = False
        
        if self.websocket:
            await self.websocket.close()
            self.websocket = None
            
        # 通知断开连接处理器
        for handler in self.disconnection_handlers:
            try:
                handler()
            except Exception as e:
                self.logger.error(f"断开连接处理器异常: {e}")
                
        self.logger.info("已断开WebSocket连接")
    
    async def send_message(self, message: Dict[str, Any]) -> bool:
        """发送消息到服务器"""
        if not self.connected or not self.websocket:
            self.logger.warning("WebSocket未连接，无法发送消息")
            return False
            
        try:
            message_str = json.dumps(message, ensure_ascii=False)
            await self.websocket.send(message_str)
            self.logger.debug(f"发送消息: {message_str}")
            return True
            
        except Exception as e:
            self.logger.error(f"发送消息失败: {e}")
            self.connected = False
            return False
    
    async def subscribe(self, symbols: list[str]) -> bool:
        """订阅股票行情"""
        message = {
            "action": "subscribe",
            "symbols": symbols
        }
        return await self.send_message(message)
    
    async def unsubscribe(self, symbols: list[str]) -> bool:
        """取消订阅股票行情"""
        message = {
            "action": "unsubscribe",
            "symbols": symbols
        }
        return await self.send_message(message)
    
    async def ping(self) -> bool:
        """发送心跳"""
        message = {"action": "ping"}
        return await self.send_message(message)
    
    def add_message_handler(self, message_type: str, handler: Callable):
        """添加消息处理器"""
        self.message_handlers[message_type] = handler
    
    def add_connection_handler(self, handler: Callable):
        """添加连接处理器"""
        self.connection_handlers.append(handler)
    
    def add_disconnection_handler(self, handler: Callable):
        """添加断开连接处理器"""
        self.disconnection_handlers.append(handler)
    
    async def _handle_message(self, message_str: str):
        """处理接收到的消息"""
        try:
            message = json.loads(message_str)
            self.logger.debug(f"收到消息: {message}")
            
            # 根据消息类型调用相应的处理器
            if "status" in message:
                # 响应消息
                message_type = "response"
            elif "symbol" in message:
                # 行情数据
                message_type = "quote"
            elif "type" in message and message["type"] == "heartbeat":
                # 心跳消息
                message_type = "heartbeat"
            else:
                message_type = "unknown"
            
            handler = self.message_handlers.get(message_type)
            if handler:
                handler(message)
            else:
                self.logger.warning(f"未找到消息类型 '{message_type}' 的处理器")
                
        except json.JSONDecodeError as e:
            self.logger.error(f"JSON解析失败: {e}")
        except Exception as e:
            self.logger.error(f"消息处理异常: {e}")
    
    async def _listen_loop(self):
        """消息监听循环"""
        while self.running and self.connected:
            try:
                if not self.websocket:
                    break
                    
                message = await asyncio.wait_for(
                    self.websocket.recv(), 
                    timeout=30.0  # 30秒超时
                )
                await self._handle_message(message)
                
            except asyncio.TimeoutError:
                # 发送心跳检测连接
                if not await self.ping():
                    self.logger.warning("心跳失败，连接可能已断开")
                    self.connected = False
                    break
                    
            except websockets.exceptions.ConnectionClosed:
                self.logger.warning("WebSocket连接已关闭")
                self.connected = False
                break
                
            except Exception as e:
                self.logger.error(f"消息监听异常: {e}")
                self.connected = False
                break
    
    async def _reconnect_loop(self):
        """重连循环"""
        while self.running:
            if not self.connected and self.reconnect_attempts < self.max_reconnect_attempts:
                self.logger.info(f"尝试重连... (第{self.reconnect_attempts + 1}次)")
                
                if await self.connect():
                    # 重连成功，启动监听循环
                    await self._listen_loop()
                else:
                    self.reconnect_attempts += 1
                    await asyncio.sleep(self.reconnect_interval)
            else:
                await asyncio.sleep(1)
    
    async def start(self):
        """启动WebSocket客户端"""
        self.running = True
        self.logger.info("启动WebSocket客户端")
        
        # 首次连接
        if await self.connect():
            # 启动监听和重连循环
            await asyncio.gather(
                self._listen_loop(),
                self._reconnect_loop()
            )
        else:
            # 连接失败，启动重连循环
            await self._reconnect_loop()
    
    async def stop(self):
        """停止WebSocket客户端"""
        self.logger.info("停止WebSocket客户端")
        await self.disconnect()


class WebSocketManager:
    """WebSocket管理器，管理与C++后端的连接"""
    
    def __init__(self, backend_uri: str = "ws://localhost:8080"):
        self.client = WebSocketClient(backend_uri)
        self.quote_cache: Dict[str, Dict] = {}
        self.subscribers: Dict[str, set] = {}  # symbol -> set of client_ids
        self.client_subscriptions: Dict[str, set] = {}  # client_id -> set of symbols
        self.logger = logging.getLogger(__name__)
        self.running = False
        self.loop_task: Optional[asyncio.Task] = None
        
        # 默认订阅的股票列表
        self.default_symbols = ["AAPL", "GOOGL", "MSFT", "AMZN", "TSLA", "META", "NVDA", "NFLX"]
        
        # 设置消息处理器
        self.client.add_message_handler("quote", self._handle_quote_update)
        self.client.add_message_handler("response", self._handle_response)
        self.client.add_message_handler("heartbeat", self._handle_heartbeat)
        
        # 添加连接成功后的处理器
        self.client.add_connection_handler(self._on_connected)
    
    def _handle_quote_update(self, message: Dict):
        """处理行情更新"""
        symbol = message.get("symbol")
        if symbol:
            self.quote_cache[symbol] = message
            self.logger.info(f"收到行情更新: {symbol} = {message.get('price')}")
    
    def _handle_response(self, message: Dict):
        """处理响应消息"""
        self.logger.info(f"收到响应: {message}")
    
    def _handle_heartbeat(self, message: Dict):
        """处理心跳消息"""
        self.logger.debug("收到心跳")
    
    def _on_connected(self):
        """连接成功后的处理"""
        self.logger.info("WebSocket连接成功，开始订阅默认股票...")
        # 创建一个任务来订阅默认股票
        asyncio.create_task(self._subscribe_default_symbols())
    
    async def _subscribe_default_symbols(self):
        """订阅默认股票"""
        try:
            # 等待一小段时间确保连接稳定
            await asyncio.sleep(0.5)
            success = await self.client.subscribe(self.default_symbols)
            if success:
                self.logger.info(f"成功订阅默认股票: {', '.join(self.default_symbols)}")
            else:
                self.logger.warning("订阅默认股票失败")
        except Exception as e:
            self.logger.error(f"订阅默认股票异常: {e}")
    
    async def start(self):
        """启动WebSocket管理器"""
        if self.running:
            return
            
        self.running = True
        self.logger.info("启动WebSocket管理器")
        
        # 在后台运行WebSocket客户端
        self.loop_task = asyncio.create_task(self.client.start())
    
    async def stop(self):
        """停止WebSocket管理器"""
        if not self.running:
            return
            
        self.running = False
        self.logger.info("停止WebSocket管理器")
        
        await self.client.stop()
        
        if self.loop_task:
            self.loop_task.cancel()
            try:
                await self.loop_task
            except asyncio.CancelledError:
                pass
    
    async def subscribe_symbols(self, symbols: list[str]) -> bool:
        """订阅股票行情"""
        return await self.client.subscribe(symbols)
    
    async def unsubscribe_symbols(self, symbols: list[str]) -> bool:
        """取消订阅股票行情"""
        return await self.client.unsubscribe(symbols)
    
    def get_quote(self, symbol: str) -> Optional[Dict]:
        """获取股票行情"""
        return self.quote_cache.get(symbol)
    
    def get_all_quotes(self) -> Dict[str, Dict]:
        """获取所有行情数据"""
        return self.quote_cache.copy()
    
    def is_connected(self) -> bool:
        """检查是否已连接"""
        return self.client.connected 